% input : P is a 2D matrix
% output : h is the matrix of H(p), element wise

function h = H_2D(p)
%   h = -p*log2(p) - (1-p)*log2(1-p);
   zero = zeros(size(p));
   one = ones(size(p));
   
   pzero = (p==zero);
   pone = (p==one);
   
   plog0 = p + pzero; % put 1 in the zero values
   plog1 = p + pone ; % put 2 in the 1 values
   
   h = -p .* log2(plog0) - (ones(size(p)) - p) .* log2(ones(size(p)) - plog1);
end